import torch
from torch import nn
from PIL import Image
from torchvision import transforms, datasets

# test_dataset = datasets.MNIST_test('./MNIST_test', train=False, download=True, transform=transforms.ToTensor())
# test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=True)


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.seq = nn.Sequential(
            nn.Conv2d(1, 28, 3, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(28, 28, 3, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(28, 56, 3, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1400, 28),
            nn.Linear(28, 10)
        )

    def forward(self, x):
        return self.seq(x)

def predict_(file_name):
    img = Image.open('../data/MNIST/MNIST_test/{}.png'.format(file_name))
    img = img.convert('L')

    img_transform = transforms.Compose([transforms.Resize((28, 28)), transforms.ToTensor()])
    img_tensor = img_transform(img)
    # transform_pil = transforms.ToPILImage()
    # pil_img = transform_pil(img_tensor)
    # pil_img.show()

    img_tensor = img_tensor.reshape(1, 1, 28, 28)
    # 训练集中的
    img_tensor = 1.0 - img_tensor
    # print(img_tensor.shape)

    with torch.no_grad():
        predict = mnist_model(img_tensor)
        predict = torch.argmax(predict, 1)

        print(int(predict.data))


if __name__ == '__main__':
    mnist_model: Net = torch.load('mnist_004.pth')
    mnist_model.eval()

    # for data_, labels in test_loader:
    #     print(data_.shape)
    #     with torch.no_grad():
    #         predict = mnist_model(data_)
    #         predict = torch.argmax(predict, 1)
    #         print(labels, predict.data)

    for i in range(10):
        print(i, ' ======predict====== ', end='')
        predict_(i)
